package mr

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"sort"
	"strconv"
	"time"
)
import "log"
import "net/rpc"
import "hash/fnv"

//
// Map functions return a slice of KeyValue.
//
type KeyValue struct {
	Key   string
	Value string
}

type WorkerInfo struct {
	Id int
}

// for sorting by key.
type ByKey []KeyValue

// for sorting by key.
func (a ByKey) Len() int           { return len(a) }
func (a ByKey) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a ByKey) Less(i, j int) bool { return a[i].Key < a[j].Key }

//
// use ihash(key) % NReduce to choose the reduce
// task number for each KeyValue emitted by Map.
//
func ihash(key string) int {
	h := fnv.New32a()
	h.Write([]byte(key))
	return int(h.Sum32() & 0x7fffffff)
}

//
// main/mrworker.go calls this function.
//
func Worker(mapf func(string, string) []KeyValue,
	reducef func(string, []string) string) {
	log.SetOutput(ioutil.Discard)

	// Your worker implementation here.
	/*
		有两种任务：Map和Reduce
		所以Worker需要根据Coordinate下发的任务类型决定具体调用map还是reduce
	*/

	// 1. 发送RPC请求，获取任务

	// 2. 判断任务类型，调用map或者reduce进行任务的处理

	// 3. 写出结果到输出文件

	var info = registerWorker()
	os.RemoveAll(strconv.Itoa(info.Id))
	log.Printf("wid: %d\n", info.Id)
	var idleCnt int = 0

	for {
		reply := TaskReply{}

		if !GetTask(&reply) {
			log.Fatalln("error doing RPC, exiting")
		}

		switch reply.Task_type {
		case Map_T:
			var outdir string
			if doMap(&info, reply, &outdir, mapf) {
				FinishMap(reply.Key, outdir)
			} else {
				log.Fatalf("error doing map with key: %s\n", reply.Key)
			}
		case Reduce_T:
			var outfile string
			if doReduce(&info, reply, &outfile, reducef) {
				FinishReduce(reply.Key, outfile)
			} else {
				log.Fatalf("error doing reduce with key: %s\n", reply.Key)
			}
		case Empty_T:
			// no task, wait
			log.Printf("No task, waiting\n")
			time.Sleep(5 * time.Second)
			idleCnt++
			if idleCnt > 10 {
				log.Println("Reached idle threshold, exiting...")
				return
			}
		}
	}
}

func registerWorker() WorkerInfo {
	args := RegisterArgs{}
	reply := RegisterReply{}
	call("Coordinator.RegisterWorker", args, &reply)
	log.Printf("got id: %d\n", reply.Wid)

	info := WorkerInfo{Id: reply.Wid}
	return info
}

func testGetTask() {
	/*
		测试获取任务功能
	*/
	args := TaskArgs{}
	reply := TaskReply{}
	call("Coordinator.GetTask", args, &reply)
	if reply.Task_type == Map_T {
		fmt.Println("map task")
	} else if reply.Task_type == Reduce_T {
		fmt.Println("reduce task")
	} else if reply.Task_type == Empty_T {
		fmt.Println("no tasks")
	}
}

func doMap(info *WorkerInfo, reply TaskReply, outdir *string, mapf func(string, string) []KeyValue) bool {
	log.Printf("running map task: %s\n", reply.Key)
	split := reply.Key
	basename := path.Base(split)

	// 读文件内容
	file, err := os.Open(split)
	if err != nil {
		log.Fatalf("cannot open %v", split)
		return false
	}
	content, err := ioutil.ReadAll(file)
	if err != nil {
		log.Fatalf("cannot read %v", split)
		return false
	}
	file.Close()

	log.Println("perform mapping")
	kva := mapf(split, string(content))

	// 创建一个新目录，将map的结果写入到其中
	*outdir = path.Join(strconv.Itoa(info.Id), basename)
	err = os.MkdirAll(*outdir, os.ModePerm)
	if err != nil {
		log.Println(err)
		log.Fatalf("failed to create dir: %s\n", *outdir)
	}

	// do partition
	var partition_files [NReduce]*os.File
	for i := 0; i < NReduce; i++ {
		part_file := path.Join(*outdir, strconv.Itoa(i))
		partition_files[i], err = os.Create(part_file)
		if err != nil {
			log.Fatalf("error creating file %s", part_file)
			return false
		}
	}

	var encs [NReduce]*json.Encoder
	for i := 0; i < NReduce; i++ {
		encs[i] = json.NewEncoder(partition_files[i])
	}

	for _, kv := range kva {
		partition := ihash(kv.Key) % NReduce
		encs[partition].Encode(kv)
	}

	for i := 0; i < NReduce; i++ {
		partition_files[i].Close()
	}
	return true
}

func doReduce(info *WorkerInfo, reply TaskReply, target *string, reducef func(string, []string) string) bool {
	/*
		读取各个目录中的第reply.Key个partition，统计其中的词频
		写入到worker私有目录
		通知coordinator任务完成
	*/
	log.Printf("running reduce task: %s\n", reply.Key)

	// 创建输出文件
	os.MkdirAll(strconv.Itoa(info.Id), os.ModePerm)
	*target = path.Join(strconv.Itoa(info.Id), "mr-out-"+reply.Key)
	out, _ := os.Create(*target)
	defer out.Close()

	// 读取所有的文件
	kva := []KeyValue{}
	for _, split := range reply.Splits {
		basename := path.Base(split)
		part_file := path.Join(MapResultDir, basename, reply.Key)
		f, err := os.Open(part_file)
		if err != nil && errors.Is(err, os.ErrNotExist) {
			if errors.Is(err, os.ErrNotExist) {
				continue
			}
			// failed to open
			log.Fatalf("failed to open file: %s", part_file)
			return false
		}
		dec := json.NewDecoder(f)
		for {
			var kv KeyValue
			if err := dec.Decode(&kv); err != nil {
				break
			}
			kva = append(kva, kv)
		}
	}

	// 排序并reduce
	sort.Sort(ByKey(kva))

	i := 0
	for i < len(kva) {
		j := i + 1
		for j < len(kva) && kva[j].Key == kva[i].Key {
			j++
		}
		values := []string{}
		for k := i; k < j; k++ {
			values = append(values, kva[k].Value)
		}
		output := reducef(kva[i].Key, values)

		// this is the correct format for each line of Reduce output.
		fmt.Fprintf(out, "%v %v\n", kva[i].Key, output)

		i = j
	}

	return true
}

// GetTask
func GetTask(reply *TaskReply) bool {
	args := new(TaskArgs)
	return call("Coordinator.GetTask", args, reply)
}

func FinishMap(key string, outdir string) {
	reply := MapTaskFinishReply{}
	args := MapTaskFinishArgs{Key: key, Output: outdir}
	call("Coordinator.FinishMapTask", args, &reply)
}

func FinishReduce(key string, outfile string) {
	reply := ReduceTaskFinishReply{}
	args := ReduceTaskFinishArgs{Key: key, Output: outfile}
	call("Coordinator.FinishReduceTask", args, &reply)
}

//
// example function to show how to make an RPC call to the coordinator.
//
// the RPC argument and reply types are defined in rpc.go.
//
func CallExample() {

	// declare an argument structure.
	args := ExampleArgs{}

	// fill in the argument(s).
	args.X = 99

	// declare a reply structure.
	reply := ExampleReply{}

	// send the RPC request, wait for the reply.
	call("Coordinator.Example", &args, &reply)

	// reply.Y should be 100.
	fmt.Printf("reply.Y %v\n", reply.Y)
}

//
// send an RPC request to the coordinator, wait for the response.
// usually returns true.
// returns false if something goes wrong.
//
func call(rpcname string, args interface{}, reply interface{}) bool {
	// c, err := rpc.DialHTTP("tcp", "127.0.0.1"+":1234")
	sockname := coordinatorSock()
	c, err := rpc.DialHTTP("unix", sockname)
	if err != nil {
		log.Fatal("dialing:", err)
	}
	defer c.Close()

	err = c.Call(rpcname, args, reply)
	if err == nil {
		return true
	}

	fmt.Println(err)
	return false
}
